module Env exposing
    ( debug
    , enter
    , gc
    , get
    , getAtom
    , global
    , globalFrameId
    , leave
    , newAtom
    , pop
    , push
    , pushRef
    , ref
    , restoreRefs
    , set
    , setAtom
    )

import Array
import Dict
import Set
import Types exposing (Env, Frame, MalExpr(..), MalFunction(..))
import Utils exposing (flip)


debug : Env -> String -> a -> a
debug env msg value =
    if env.debug then
        Debug.log msg value

    else
        value


globalFrameId : Int
globalFrameId =
    0


defaultGcInterval : Int
defaultGcInterval =
    10


global : Env
global =
    { frames = Dict.singleton globalFrameId (emptyFrame Nothing Nothing)
    , nextFrameId = globalFrameId + 1
    , currentFrameId = globalFrameId
    , atoms = Dict.empty
    , nextAtomId = 0
    , debug = False
    , gcInterval = defaultGcInterval
    , gcCounter = 0
    , stack = []
    , keepFrames = []
    }


getFrame : Env -> Int -> Frame
getFrame env frameId =
    case Dict.get frameId env.frames of
        Just frame ->
            frame

        Nothing ->
            Debug.todo <| "frame #" ++ String.fromInt frameId ++ " not found"


emptyFrame : Maybe Int -> Maybe Int -> Frame
emptyFrame outerId exitId =
    { outerId = outerId
    , exitId = exitId
    , data = Dict.empty
    , refCnt = 1
    }


set : String -> MalExpr -> Env -> Env
set name expr env =
    let
        frameId =
            env.currentFrameId

        updateFrame =
            Maybe.map
                (\frame ->
                    { frame | data = Dict.insert name expr frame.data }
                )

        newFrames =
            Dict.update frameId updateFrame env.frames
    in
    { env | frames = newFrames }


get : String -> Env -> Result String MalExpr
get name env =
    let
        go frameId =
            let
                frame =
                    getFrame env frameId
            in
            case Dict.get name frame.data of
                Just value ->
                    Ok value

                Nothing ->
                    frame.outerId
                        |> Maybe.map go
                        |> Maybe.withDefault (Err <| "'" ++ name ++ "' not found")
    in
    go env.currentFrameId


newAtom : MalExpr -> Env -> ( Env, Int )
newAtom value env =
    let
        atomId =
            env.nextAtomId

        newEnv =
            { env
                | atoms = Dict.insert atomId value env.atoms
                , nextAtomId = atomId + 1
            }
    in
    ( newEnv, atomId )


getAtom : Int -> Env -> MalExpr
getAtom atomId env =
    case Dict.get atomId env.atoms of
        Just value ->
            value

        Nothing ->
            Debug.todo <| "atom " ++ String.fromInt atomId ++ " not found"


setAtom : Int -> MalExpr -> Env -> Env
setAtom atomId value env =
    { env
        | atoms = Dict.insert atomId value env.atoms
    }


push : Env -> Env
push env =
    let
        frameId =
            env.nextFrameId

        newFrame =
            emptyFrame (Just env.currentFrameId) Nothing

        bogus =
            debug env "push" frameId
    in
    { env
        | currentFrameId = frameId
        , frames = Dict.insert frameId newFrame env.frames
        , nextFrameId = env.nextFrameId + 1
    }


pop : Env -> Env
pop env =
    let
        frameId =
            env.currentFrameId

        frame =
            getFrame env frameId

        bogus =
            debug env "pop" frameId
    in
    case frame.outerId of
        Just outerId ->
            { env
                | currentFrameId = outerId
                , frames = Dict.update frameId free env.frames
            }

        _ ->
            Debug.todo "tried to pop global frame"


setBinds : List ( String, MalExpr ) -> Frame -> Frame
setBinds binds frame =
    case binds of
        [] ->
            frame

        ( name, expr ) :: rest ->
            setBinds rest
                { frame | data = Dict.insert name expr frame.data }


{-| Enter a new frame with a set of binds
-}
enter : Int -> List ( String, MalExpr ) -> Env -> Env
enter outerId binds env =
    let
        frameId =
            debug env "enter #" env.nextFrameId

        exitId =
            env.currentFrameId

        newFrame =
            setBinds binds (emptyFrame (Just outerId) (Just exitId))
    in
    { env
        | currentFrameId = frameId
        , frames = Dict.insert frameId newFrame env.frames
        , nextFrameId = env.nextFrameId + 1
    }


leave : Env -> Env
leave env =
    let
        frameId =
            debug env "leave #" env.currentFrameId

        frame =
            getFrame env frameId

        exitId =
            case frame.exitId of
                Just exitId2 ->
                    exitId2

                Nothing ->
                    Debug.todo <|
                        "frame #"
                            ++ String.fromInt frameId
                            ++ " doesn't have an exitId"
    in
    { env
        | currentFrameId = exitId
        , frames =
            env.frames
                |> Dict.insert frameId { frame | exitId = Nothing }
                |> Dict.update frameId free
    }


{-| Increase refCnt for the current frame,
and all it's parent frames.
-}
ref : Env -> Env
ref originalEnv =
    let
        go frameId env =
            let
                frame =
                    getFrame env frameId

                newFrame =
                    { frame | refCnt = frame.refCnt + 1 }

                newEnv =
                    { env | frames = Dict.insert frameId newFrame env.frames }
            in
            case frame.outerId of
                Just outerId ->
                    go outerId newEnv

                Nothing ->
                    newEnv

        newEnv2 =
            go originalEnv.currentFrameId originalEnv
    in
    { newEnv2 | gcCounter = newEnv2.gcCounter + 1 }


free : Maybe Frame -> Maybe Frame
free =
    Maybe.andThen
        (\frame ->
            if frame.refCnt == 1 then
                Nothing

            else
                Just { frame | refCnt = frame.refCnt - 1 }
        )


pushRef : MalExpr -> Env -> Env
pushRef ref_arg env =
    { env | stack = ref_arg :: env.stack }


restoreRefs : List MalExpr -> Env -> Env
restoreRefs refs env =
    { env | stack = refs }


{-| Given an Env see which frames are not reachable from the
global frame, or from the current expression.

Return a new Env with the unreachable frames removed.

-}
gc : MalExpr -> Env -> Env
gc expr env =
    let
        countList acc =
            List.foldl countExpr acc

        countFrame { data } acc =
            data |> Dict.values |> countList acc

        recur frameId acc =
            if not (Set.member frameId acc) then
                let
                    frame =
                        getFrame env frameId

                    newAcc =
                        Set.insert frameId acc
                in
                countFrame frame newAcc

            else
                acc

        countBound bound acc =
            bound
                |> List.map Tuple.second
                |> countList acc

        countExpr expr_arg acc =
            case expr_arg of
                MalFunction (UserFunc { frameId }) ->
                    recur frameId acc

                MalApply { frameId, bound } ->
                    recur frameId acc
                        |> countBound bound

                MalList _ list ->
                    countList acc list

                MalVector _ vec ->
                    countList acc (Array.toList vec)

                MalMap _ map ->
                    countList acc (Dict.values map)

                MalAtom atomId ->
                    let
                        value =
                            getAtom atomId env
                    in
                    countExpr value acc

                _ ->
                    acc

        initSet =
            Set.fromList
                ([ globalFrameId, env.currentFrameId ]
                    ++ env.keepFrames
                )

        countFrames frames acc =
            Set.toList frames
                |> List.map (getFrame env)
                |> List.foldl countFrame acc

        expand frameId frame fn acc =
            case fn frame of
                Nothing ->
                    acc

                Just parentId ->
                    Set.insert parentId acc

        expandBoth frameId =
            let
                frame =
                    getFrame env frameId
            in
            expand frameId frame .outerId
                >> expand frameId frame .exitId

        expandParents frames =
            Set.foldl expandBoth frames frames

        loop acc =
            let
                newAcc =
                    expandParents acc

                newParents =
                    Set.diff newAcc acc
            in
            if Set.isEmpty newParents then
                newAcc

            else
                loop <| countFrames newParents newAcc

        makeNewEnv newFrames =
            { env
                | frames = newFrames
                , gcCounter = 0
            }

        keepFilter keep frameId _ =
            Set.member frameId keep

        filterFrames frames keep =
            Dict.filter (keepFilter keep) frames
    in
    countFrames initSet initSet
        |> countExpr expr
        |> flip countList env.stack
        |> loop
        |> filterFrames env.frames
        |> makeNewEnv
